import numpy as np
from util import createXY
from sklearn.model_selection import train_test_split
import logging
from time import time
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import (
    RandomForestClassifier,
    VotingClassifier,
    BaggingClassifier,
    AdaBoostClassifier,
    GradientBoostingClassifier
)
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score

# 忽略警告信息
import warnings
warnings.filterwarnings('ignore')

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def train_and_evaluate_model(model, X_train, X_test, y_train, y_test, model_name):
    """训练并评估模型，返回训练时间、预测时间和准确率"""
    train_start = time()
    model.fit(X_train, y_train)
    train_time = time() - train_start
    pred_start = time()
    y_pred = model.predict(X_test)
    pred_time = time() - pred_start
    accuracy = accuracy_score(y_test, y_pred)
    logging.info(f"{model_name} - 训练时间: {train_time:.4f}s, 预测时间: {pred_time:.4f}s, 准确率: {accuracy:.4f}")
    return train_time, pred_time, accuracy

def main():
    # 加载数据
    X, y = createXY(train_folder="../data/train", dest_folder=".", method='flat')
    X = np.array(X, dtype='float32')
    y = np.array(y)

    # 数据预处理
    logging.info("数据加载和预处理完成。")

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")

    # 初始化模型
    models = {
        'logistic_regression': LogisticRegression(max_iter=1000),
        'random_forest': RandomForestClassifier(n_estimators=100, random_state=2023),
        'svm': SVC(kernel='rbf', probability=True),
        'hard_voting': VotingClassifier(
            estimators=[
                ('lr', LogisticRegression(max_iter=1000)),
                ('rf', RandomForestClassifier(n_estimators=100, random_state=2023)),
                ('svm', SVC(kernel='rbf', probability=True))
            ],
            voting='hard'
        ),
        'soft_voting': VotingClassifier(
            estimators=[
                ('lr', LogisticRegression(max_iter=1000)),
                ('rf', RandomForestClassifier(n_estimators=100, random_state=2023)),
                ('svm', SVC(kernel='rbf', probability=True))
            ],
            voting='soft'
        ),
        'bagging': BaggingClassifier(
            base_estimator=LogisticRegression(max_iter=1000),
            n_estimators=10,
            random_state=2023
        ),
        'pasting': BaggingClassifier(
            base_estimator=LogisticRegression(max_iter=1000),
            n_estimators=10,
            bootstrap=False,
            random_state=2023
        ),
        'adaboost': AdaBoostClassifier(
            base_estimator=LogisticRegression(max_iter=1000),
            n_estimators=50,
            random_state=2023
        ),
        'gradient_boosting': GradientBoostingClassifier(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=3,
            random_state=2023
        )
    }

    # 存储结果
    results = []

    # 训练和评估每个模型
    for model_name, model in models.items():
        train_time, pred_time, accuracy = train_and_evaluate_model(
            model, X_train, X_test, y_train, y_test, model_name
        )
        results.append({
            'Classifier': model_name,
            'Training Time (s)': train_time,
            'Prediction Time (s)': pred_time,
            'Accuracy': accuracy
        })

    # 创建结果数据框并打印
    results_df = pd.DataFrame(results)
    logging.info("\n模型评估结果：")
    logging.info(results_df.to_string(index=False))

if __name__ == '__main__':
    main()